#include"mymath.h"
#include<iostream>
#include <string.h>
const double pi = acos(-1);
MYComplex operator + (MYComplex a, MYComplex b) { return MYComplex(a.r + b.r, a.i + b.i); }
MYComplex operator - (MYComplex a, MYComplex b) { return MYComplex(a.r - b.r, a.i - b.i); }
MYComplex operator * (MYComplex a, MYComplex b) { return MYComplex(a.r * b.r - a.i * b.i, a.r * b.i + a.i * b.r); }
MYComplex operator / (MYComplex a, MYComplex b) { return MYComplex((a.r * b.r + a.i * b.i) / (b.r * b.r + b.i * b.i), (a.i * b.r - a.r * b.i) / (b.r * b.r + b.i * b.i));}
void FFT(MYComplex* a, int opt, int n) {
	int len = 0, lim = 1;
	while (lim < n) lim <<= 1, len++;
	int* rev = (int*)malloc(sizeof(int) * lim);
	memset(rev, 0, sizeof(rev));

	for (int i = 0; i < lim; ++i) rev[i] = (rev[i >> 1] >> 1) | ((i & 1) << (len - 1));
	for (int i = 0; i < lim; ++i) if (i < rev[i]) std::swap(a[i], a[rev[i]]);
	for (int dep = 1; dep <= log2(lim); ++dep) {
		int m = 1 << (dep);
		MYComplex wn = MYComplex(cos(2.0 * pi / m), opt * (-1) * sin(2.0 * pi / m));
		for (int k = 0; k < lim; k += m) {
			MYComplex w = MYComplex(1, 0);
			for (int j = 0; j < m / 2; ++j) {
				MYComplex t = w * a[k + j + m / 2];
				MYComplex u = a[k + j];
				a[k + j] = u + t;
				a[k + j + m / 2] = u - t;
				w = w * wn;
			}
		}
	}
	if (opt == -1) for (int i = 0; i < lim; ++i) a[i].r /= lim, a[i].i /= lim;
	free(rev);
}
// void FFTshift(std::vector<cv::Mat>& src) {
// 	int temp = src[0].cols / 2;
// 	for (int i = 0; i < src[0].rows; i++) {
// 		for (int j = 0; j < temp; j++) {
// 			std::swap(src[0].ptr<double>(i)[j], src[0].ptr<double>(i)[j + temp]);
// 			std::swap(src[1].ptr<double>(i)[j], src[1].ptr<double>(i)[j + temp]);
// 		}
// 	}
// 	temp = src[0].rows / 2;
// 	for (int i = 0; i < temp; i++) {
// 		for (int j = 0; j < src[0].cols; j++) {
// 			std::swap(src[0].ptr<double>(i)[j], src[0].ptr<double>(i + temp)[j]);
// 			std::swap(src[1].ptr<double>(i)[j], src[1].ptr<double>(i + temp)[j]);
// 		}
// 	}
// }
void FFTshift(std::vector<std::vector<std::vector<float>>>& src) {
	int klen = src.size(), ilen = src[0].size(), jlen = src[0][0].size();
	int temp = jlen>>1;
	for(int k = 0 ; k < klen ; k ++){
		for (int i = 0; i < ilen; i++) {
			for (int j = 0; j < temp; j++) {
				std::swap(src[k][i][j],src[k][i][j+temp]);
			}
		}
	}
	temp = ilen>>1;
	for(int k = 0 ; k < klen ; k ++){
		for (int i = 0; i < temp; i++) {
			for (int j = 0; j < jlen; j++) {
				std::swap(src[k][i][j],src[k][i+temp][j]);
			}
		}
	}
}
// void FFT2D(int height, int width, int opt, std::vector<Mat>& src, std::vector<Mat>& dst) {
// 	int i, j;
// 	int paddingh = 1, paddingw = 1, bitlenh = 0, bitlenw = 0;
// 	while (paddingh < height) paddingh <<= 1, bitlenh++;
// 	while (paddingw < width) paddingw <<= 1, bitlenw++;
// 	if(debug) printf("FF2D(%d): clear dst start\n", opt);
// 	dst.clear();
// 	dst.push_back(cv::Mat::zeros(paddingh, paddingw, CV_64FC1));
// 	dst.push_back(cv::Mat::zeros(paddingh, paddingw, CV_64FC1));
// 	if (debug) printf("FF2D(%d): clear dst done\n", opt);
// 	//for (i = 0; i < src[0].rows; i++) {
// //	for (j = 0; j < src[0].cols; j++) {
// //		dst[0].ptr<double>(i)[j] = src[0].ptr<double>(i)[j];
// //		dst[1].ptr<double>(i)[j] = src[1].ptr<double>(i)[j];
// //	}
// //}
// 	for (i = 0; i < paddingh; i++) {
// 		if (i >= 0 && i < src[0].rows) {
// 			for (j = 0; j < src[0].cols; j++) {
// 				dst[0].ptr<double>(i)[j] = src[0].ptr<double>(i)[j];
// 				dst[1].ptr<double>(i)[j] = src[1].ptr<double>(i)[j];
// 			}
// 			for (j = src[0].cols; j < paddingw; j++) {
// 				dst[0].ptr<double>(i)[j] = src[0].ptr<double>(i)[2 * src[0].cols - j - 1];
// 				dst[1].ptr<double>(i)[j] = src[1].ptr<double>(i)[2 * src[0].cols - j - 1];
// 			}
// 		}
// 		else if(i >= src[0].rows && i < paddingh) {
// 			for (j = 0; j < paddingw; j++) {
// 				dst[0].ptr<double>(i)[j] = dst[0].ptr<double>(2 * src[0].rows - i - 1)[j];
// 				dst[1].ptr<double>(i)[j] = dst[1].ptr<double>(2 * src[0].rows - i - 1)[j];
// 			}
// 		}
// 	}
// 	if (debug) printf("FF2D(%d): padding done\n", opt);
// 	if (opt == -1) FFTshift(dst);
// 	if (debug) printf("FF2D(%d): IFFTshift done\n", opt);
// 	int maxx = paddingh > paddingw ? paddingh : paddingw;
// 	MYComplex* A = (MYComplex*)malloc(sizeof(MYComplex) * maxx);
// 	for (i = 0; i < paddingh; i++) {
// 		for (j = 0; j < paddingw; j++) {
// 			A[j].r = dst[0].ptr<double>(i)[j];
// 			A[j].i = dst[1].ptr<double>(i)[j];
// 			//printf("%.4lf + (%.4lf)i\n", A[j].r, A[j].i);
// 		}
// 		FFT(A, opt, paddingw);
// 		for (j = 0; j < paddingw; j++)
// 		{
// 			dst[0].ptr<double>(i)[j] = A[j].r;
// 			dst[1].ptr<double>(i)[j] = A[j].i;
// 		}
// 	}
// 	if (debug) printf("FF2D(%d) done\n", opt);
// 	for (i = 0; i < paddingw; i++)
// 	{
// 		for (j = 0; j < paddingh; j++)
// 		{
// 			A[j].r = dst[0].ptr<double>(j)[i];
// 			A[j].i = dst[1].ptr<double>(j)[i];
// 		}
// 		FFT(A, opt, paddingh);
// 		for (j = 0; j < paddingh; j++)
// 		{
// 			dst[0].ptr<double>(j)[i] = A[j].r;
// 			dst[1].ptr<double>(j)[i] = A[j].i;
// 		}
// 	}
// 	if (debug) printf("FF2D(%d) done\n", opt);
// 	if (opt == 1) FFTshift(dst);
// 	free(A);
// }
void FFT2D(int height, int width, int opt, std::vector<std::vector<std::vector<float>>>& src, std::vector<std::vector<std::vector<float>>>& DstRe) {
	int i, j;
	if(debug) printf("FF2D(%d): init parameters\n", opt);
	int klen = src.size() >> 1, ilen = src[0].size(), jlen = src[0][0].size();
	int paddingh = 1, paddingw = 1, bitlenh = 0, bitlenw = 0;
	while (paddingh < height) paddingh <<= 1, bitlenh++;
	while (paddingw < width) paddingw <<= 1, bitlenw++;
	
	if(debug) printf("FF2D(%d): init parameters Done(paddingh:%d paddingw:%d klen:%d ilen:%d jlen:%d)\n", opt,paddingh,paddingw,klen,ilen,jlen);

	std::vector<std::vector<std::vector<float>>>DstIm;
	std::vector<std::vector<float>>DstReMap,DstImMap;
	std::vector<float>DstReLine,DstImLine;

	if(debug) printf("FF2D(%d): clear dst start\n", opt);
	DstRe.clear();
	DstIm.clear();
	if (debug) printf("FF2D(%d): clear dst done\n", opt);
	
	if(debug) printf("FF2D(%d):mirror padding start\n", opt);
	for(int k = 0 ; k < klen ; k ++){
		DstReMap.clear();
		DstImMap.clear();
		for(i = 0 ; i < paddingh ; i ++){
			DstReLine.clear();
			DstImLine.clear();

			if (i >= 0 && i < ilen) {
				for(j = 0 ; j < jlen ; j ++){
					DstReLine.push_back(src[k][i][j]);
					DstImLine.push_back(src[k+klen][i][j]);
				}
				for (j = jlen; j < paddingw; j++) {
					DstReLine.push_back(src[k][i][(jlen << 1) - j - 1]);
					DstImLine.push_back(src[k+klen][i][(jlen << 1) - j - 1]);
				}
			}
			else if(i >= ilen && i < paddingh) {
				for (j = 0; j < paddingw; j++) {
					DstReLine.push_back(src[k][(ilen << 1) - i - 1][j]);
					DstImLine.push_back(src[k+klen][(ilen << 1) - i - 1][j]);
				}
			}
			DstReMap.push_back(DstReLine);
			DstImMap.push_back(DstImLine);
		}
		DstRe.push_back(DstReMap);
		DstIm.push_back(DstImMap);
	}
	if(debug) printf("FF2D(%d):mirror padding Done\n", opt);

	if (debug) printf("FF2D(%d): IFFTshift Start\n", opt);
	if (opt == -1){FFTshift(DstRe);FFTshift(DstIm);}
	if (debug) printf("FF2D(%d): IFFTshift done\n", opt);
	int maxx = paddingh > paddingw ? paddingh : paddingw;
	MYComplex* A = (MYComplex*)malloc(sizeof(MYComplex) * maxx);

	if (debug) printf("FF2D(%d): calculating Start(maxx:%d Size of DstRe:(%d,%d,%d) Size of DstIm:(%d,%d,%d)\n", 
						opt,maxx,DstRe.size(),DstRe[0].size(),DstRe[0][0].size(),DstIm.size(),DstIm[0].size(),DstIm[0][0].size());
	for(int k = 0 ; k < klen ; k ++){

		for (i = 0; i < paddingh; i++) {
			for (j = 0; j < paddingw; j++) {
				A[j].r = DstRe[k][i][j];
				A[j].i = DstIm[k][i][j];
				//printf("%.4lf + (%.4lf)i\n", A[j].r, A[j].i);
			}

			FFT(A, opt, paddingw);
			for (j = 0; j < paddingw; j++)
			{
				DstRe[k][i][j] = A[j].r;
				DstIm[k][i][j] = A[j].i;
			}
		}
		if (debug) printf("FF2D(%d): %d-th half done\n", opt,k);
		for (i = 0; i < paddingw; i++)
		{
			for (j = 0; j < paddingh; j++)
			{
				A[j].r = DstRe[k][j][i];
				A[j].i = DstIm[k][j][i];
			}


			FFT(A, opt, paddingh);
			for (j = 0; j < paddingh; j++)
			{
				DstRe[k][j][i] = A[j].r;
				DstIm[k][j][i] = A[j].i;
			}
		}
		if (debug) printf("FF2D(%d): %d-th all done\n", opt,k);
	}
	if (debug) printf("FF2D(%d): calculating Done\n", opt);
	if (debug) printf("FF2D(%d): FFTSHIFT Start\n", opt);
	if (opt == 1){FFTshift(DstRe);FFTshift(DstIm);}
	if (debug) printf("FF2D(%d): FFTSHIFT Done\n", opt);
	if (debug) printf("FF2D(%d): datamerge Start\n", opt);
	DstRe.insert(DstRe.end(),DstIm.begin(),DstIm.end());
	if (debug) printf("FF2D(%d): datamerge Done\n", opt);
	free(A);
}